import grpc
import calligraphy_pb2_grpc,calligraphy_pb2

import os
from PIL import Image
import torchvision.transforms as transforms
import glob
import matplotlib.pyplot as plt
import numpy as np
import time
import torch
import base64
import cv2
from io import BytesIO,StringIO

preprocessing_train = transforms.Compose([
    transforms.RandomResizedCrop(224),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
  transforms.Normalize(mean=[0.485, 0.456, 0.406],
                       std=[0.229, 0.224, 0.225]),
])

preprocessing_val = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225]),
])



def run(image, height, width, gt_label):
  with grpc.insecure_channel('localhost:2534') as channel:
  # with grpc.insecure_channel('207.246.117.29:8080') as channel:
    stub = calligraphy_pb2_grpc.NetStub(channel=channel)
    request = calligraphy_pb2.ImageMatrix(image=image, height=height, width=width)

    start = time.time()
    response = stub.ClassifyImage(request)
    end = time.time()
  print(" =======>   image category: " + str(int(response.category)),
        "  name: ", id_to_label[int(response.category)],
        "  gt_category: ", gt_label,
        "  gt_namge: ", id_to_label[gt_label], "\n",
        "           forward time : ", response.run_time,
        "  running time : ", end - start)

if __name__ == '__main__':
  a = os.listdir("/media/retoo/RetooDisk1/wanghui/Data/ILSVR2012C/ILSVRC2012_img_val/")
  category = open("../imagenet1000_clsidx_to_labels.txt","r").readlines()
  id_to_label = dict()
  for line in category:
    line = line.split(":")
    id_to_label[int(line[0])] = line[1].strip()

  root = "/media/retoo/RetooDisk1/wanghui/Data/ILSVR2012C/ILSVRC2012_img_val/"
  files = open("/media/retoo/RetooDisk1/wanghui/Data/ILSVR2012C/ILSVRC2012_val.txt", "r")
  for line in files.readlines():
    line = line.strip().split(" ")
    gt_label = int(line[1])
    p = os.path.join(root, line[0].split("/")[-1])
    src = Image.open(p).convert('RGB')
    width, height = src.size
    # image = np.asarray(src)
    # string = base64.b64encode(image)

    # _, ext = os.path.splitext(os.path.basename(p))
    # img_data = cv2.imread(p)[:, :, ::-1]  # read image and convert to RGB channel
    # _, img_encode = cv2.imencode(ext, img_data)
    # data_encode = np.array(img_encode)
    # str_encode = data_encode.tobytes()

    encode_image = base64.b64encode(open(p, "rb").read())
    run(encode_image, height, width, gt_label)

    # img = preprocessing_val(src)
    # src = src.resize((224,224))
    #
    # src = np.asarray(src)
    # img = torch.from_numpy(src).permute((2, 0, 1)).contiguous()
    # img = img.unsqueeze(dim=0)
    # run(img.view(-1),gt_label)
    # plt.imshow(np.asarray(src))
    # plt.show()


# docker build --pull -t dl-pytorch-serving -f ./torch_grpc_python_serving/docker_without_proxy/pytorch_gpu_Dockerfile .
# docker build --pull -t dl-pytorch-serving -f ./torch_grpc_python_serving/docker_without_proxy/pytorch_cpu_Dockerfile .
# docker run -p 45.77.51.15:2500:2500 --name=pytorch_service -d -it dl-pytorch-serving:latest python3 ./server.py
# docker run -p 45.32.188.63:2500:2500 --name=pytorch_service_grpc -it dl-pytorch-serving:latest python3 ./server.py
# docker run -p 207.246.117.29:2501:2501 --name=pytorch_service_flask -it dl-pytorch-serving:latest python3 ./server_restapi.py

# gpu version
# docker run -p 8080:8080 --gpus all --name=pytorch_service_gpu -it yinchen0243/resnet-pytorch-serving python3 ./server.py
# docker run -p 8080:8080  --name=pytorch_service_cpu -it yinchen0243/resnet-pytorch-serving python3 ./server.py